# 将优化策略整合到RAG链中

# 优质的检索 + 严格的Prompt工程 + 答案验证/事实核查 + 模型选择与微调
# 我们将 EnsembleRetriever 作为 MultiQueryRetriever 的基础。
# 然后将 MultiQueryRetriever 的输出，再送入 ContextualCompressionRetriever 中，使用 FlashrankRerank (重排序) 进行最后的精简和排序。
# 这个 final_optimized_retriever 就是一个实现了“先扩展查询 -> 再混合搜索 -> 最后重排序结果”的检索器。
# 最终，我们将这个优化后的检索器整合到标准的 create_retrieval_chain 中，构建出一个非常强大的RAG应用。

# pip install rank_bm25 -i https://pypi.tuna.tsinghua.edu.cn/simple/

# from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
# from langchain.retrievers.document_compressors import FlashrankRerank
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers import EnsembleRetriever
# from langchain_chroma import Chroma
# from langchain_community.vectorstores import Chroma
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import MultiQueryRetriever
from langchain_community.document_compressors import FlashrankRerank
from langchain_core.documents import Document
from langchain_community.vectorstores import FAISS

from langchain_ollama import OllamaEmbeddings
embeddings = OllamaEmbeddings(model="bge-m3:latest")

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

split_documents = [    
    Document(page_content="LCEL，全称LangChain Expression Language，它通过操作符重载（如`|`符号）提供了一种声明式的、流畅的方式来构建AI链。它的关键优势包括：开箱即用的流式处理、异步和并行执行能力，以及对整个链的生命周期管理（如日志、调试）提供了极大的便利。", metadata={"doc_id": "lcel_intro"}),
    Document(page_content="混合搜索（Hybrid Search）结合了传统关键词搜索（如BM25）和现代向量搜索的优点。关键词搜索能精确匹配术语和缩写，而向量搜索擅长理解语义和意图。二者结合能显著提升检索的鲁棒性和准确性。", metadata={"doc_id": "hybrid_search_intro"}),
]

vectorstore = FAISS.from_documents(split_documents, embeddings)

# --- 1. 构建多阶段优化检索器 ---
# 基础检索器: 混合搜索
bm25_retriever = BM25Retriever.from_documents(split_documents)
bm25_retriever.k = 5

vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

ensemble_retriever = EnsembleRetriever(    
    retrievers=[bm25_retriever, vector_retriever],    
    weights=[0.5, 0.5])

# 查询扩展: 将混合搜索作为基础
multi_query_hybrid_retriever = MultiQueryRetriever.from_llm(    
    retriever=ensemble_retriever,    
    llm=llm)

# 重排序: 将查询扩展和混合搜索的结果进行精排
rerank_compressor = FlashrankRerank(    
    model="miniReranker_arabic_v1",    
    top_n=3)

# 指定本地模型路径
# rerank_compressor = FlashrankRerank(
#     model="path/to/your/local/model.onnx",  # 本地模型路径
#     top_n=3,
#     model_file="local"  # 关键参数，声明使用本地模型
# )

# 创建一个 ContextualCompressionRetriever, 使用 Flashrank 作为压缩器
final_optimized_retriever = ContextualCompressionRetriever(    
    base_compressor=rerank_compressor,    
    base_retriever=vectorstore.as_retriever(search_kwargs={"k": 5}) # 先检索5个再重排
)

# --- 2. 构建最终的RAG链 ---# 
# a. 定义RAG的Prompt
optimized_rag_prompt = ChatPromptTemplate.from_messages([    
    ("system", "你是一名专业的知识库助手。请根据提供的上下文**简洁明了**地回答以下问题。\n**如果上下文没有足够信息，请明确说明你不知道，不要凭空捏造。**\n\n上下文:\n{context}"),    
    ("user", "{input}")
])

# 创建文档处理链
optimized_document_chain = create_stuff_documents_chain(llm, optimized_rag_prompt)

# 创建完整的检索链
optimized_retrieval_chain = create_retrieval_chain(    
    final_optimized_retriever,    
    optimized_document_chain)

# --- 3. 调用测试 ---
print("\n--- 完整优化后的RAG链示例 ---")
query_final = "LangChain的优势是什么？"
response = optimized_retrieval_chain.invoke({"input": query_final})
print(f"用户: {query_final}")
print(f"AI: {response['answer']}")

# 查看检索到的上下文，验证其高质量
print("\n--- 检索到的上下文(Context) ---")
for i, doc in enumerate(response['context']):    
    print(f"文档 {i+1} (来源: {doc.metadata.get('source', 'N/A')}, 分数: {doc.metadata.get('relevance_score', 'N/A')}):")    
    print(doc.page_content)    
    print("-" * 20)
